#include <iostream>
#include <optional>
#include <fstream>
#include <filesystem>
#include <cnpy.h>
#include <nlohmann/json.hpp>
#include <glm/glm.hpp>
#include <fmt/format.h>

#include <eustis/AppFrame.h>
#include <eustis/Entity.h>
#include <eustis/Camera.h>
#include <eustis/CameraController.h>
#include <eustis/PrimitiveHelper.h>
#include <eustis/Renderable.h>
#include <eustis/VertexType.h>
#include <eustis/SceneNode.h>
#include <eustis/GLEffect.h>
#include <eustis/Material.h>
#include <eustis/ShaderStorageBuffer.h>
#include <eustis/TextureHelper.h>
#include <eustis/TextureTypes.h>
#include <eustis/Texture.h>
#include <eustis/ImgppTextureSrc.h>
#include <eustis/Engine.h>
#include <eustis/glutility.h>

using namespace eustis;

namespace {

void print_npy_array_info(const cnpy::NpyArray &array) {
  std::cout << array.num_vals << " * " << array.word_size << std::endl;
  for (auto v : array.shape) {
    std::cout << v << ", ";
  }
  std::cout << std::endl << std::endl;;
}

class TextureCubeMap : public InternalTexture {
public:
  TextureCubeMap() {}
  virtual ~TextureCubeMap() {}
  bool Create(const TextureDesc &desc) override {
    if (desc.type != GL_TEXTURE_CUBE_MAP) {
      return false;
    }
    if (desc.compressed) {
      return false;
    }
    if (desc.samples > 1 && desc.levels > 1) {
      return false;
    }

    desc_ = desc;
    glBindTexture(desc_.type, handle_);

    int actual_levels = desc_.levels;
    if (actual_levels == 0) {  // manually calculate max level
      actual_levels = (int)std::floor(std::log2((float)std::max(desc_.width, desc_.height))) + 1;
    }

    //setup sampling parameters
    glTexParameteri(desc_.type, GL_TEXTURE_MAG_FILTER, gl::MapFilterMode(desc_.mag_filter));
    glTexParameteri(desc_.type, GL_TEXTURE_MIN_FILTER, gl::MapFilterMode(desc_.min_filter));
    glTexParameteri(desc_.type, GL_TEXTURE_WRAP_S, gl::MapWrapMode(desc_.wrap_s));
    glTexParameteri(desc_.type, GL_TEXTURE_WRAP_T, gl::MapWrapMode(desc_.wrap_t));
    glTexParameteri(desc_.type, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);
    glTexParameteri(desc_.type, GL_TEXTURE_MAX_LEVEL, actual_levels - 1);
    glTexParameteri(desc_.type, GL_TEXTURE_COMPARE_MODE,
      desc_.depth_compare ? GL_COMPARE_REF_TO_TEXTURE : GL_NONE);
    glTexParameteri(desc_.type, GL_TEXTURE_COMPARE_FUNC, desc_.depth_func);

    glm::ivec3 dims = desc_.src_data->Dimensions(0);
    int alignment = desc_.src_data->Alignment(0);
    glPixelStorei(GL_UNPACK_ALIGNMENT, alignment);
    for (auto i = 0; i < 6; ++i) {
      glTexImage2D(GL_TEXTURE_CUBE_MAP_POSITIVE_X + i, 0, desc_.internal_format,
        dims[0], dims[1], 0, desc_.format, desc_.data_type,
        desc_.src_data->Data(i, 0, 0));
    }
    glPixelStorei(GL_UNPACK_ALIGNMENT, 4);  // resotre pack alignment
    glGenerateMipmap(desc_.type);

    glBindTexture(desc_.type, 0);
    desc_.src_data.reset();
    CHKGLERR_RET

    return true;
  }
  bool UpdateTexture(const SubTextureDesc &desc, const void *data) override {
    return false;
  }
};

class MsiEntity : public Entity {
public:
  MsiEntity(const std::string &data_folder, bool flip_y) :
    flip_y_(flip_y) {
    folder_ = std::filesystem::path(data_folder);
    if (!std::filesystem::exists(folder_ / "params.json")) {
      throw std::runtime_error("params.json does not exist");
    }
    std::ifstream f(folder_ / "params.json");
    f >> params_;
  }

  void Start() override {
    Engine &en = Engine::Instance();
    auto scene = en.Scene();

    // create camera
    auto camera = std::make_shared<Camera>();
    camera->SetProjParams(glm::radians(60.0f), 0.1, 100);
    en.CameraMgr().Add("camera", camera);
    en.ChangeCamera("camera", false);
    auto camera_node = std::make_shared<SceneNode>("camera");
    glm::vec3 up_dir = flip_y_?glm::vec3(0, -1, 0):glm::vec3(0, 1, 0);
    camera_node->SetRbt(math::Rbt::LookAt(
      glm::vec3(0, 0, 0), glm::vec3(0, 0, -1), up_dir));
    camera_node->Cameras().push_back(camera);
    SceneNode::LinkTo(camera_node, scene);
    cam_ctrl_.reset(new FPSController);
    cam_ctrl_->BindToNode("camera");
    cam_ctrl_->SetUpDir(up_dir);
    cam_ctrl_->Activate();

    // params from config
    int msi_layers = params_["msi_layers"].get<int>();
    int total_layers = params_["total_layers"].get<int>();
    int msi_sublayers = total_layers / msi_layers;
    int msi_w = params_["msi_w"].get<int>();
    int msi_h = params_["msi_h"].get<int>();
    float theta_half_range = glm::radians(
      (float)params_["theta_half_degree_range"].get<int>());
    float phi_half_range = glm::radians(
      (float)params_["phi_half_degree_range"].get<int>());
    std::vector<float> radii = params_["radii"].get<std::vector<float>>();
    int basis_out = params_["basis_out"].get<int>();

    // create total_layers planes to draw each radius
    for (int i = 0; i < total_layers; ++i) {
      auto renderable = std::make_shared<Renderable>();
      auto vertex_array = std::make_shared<vertex_type::VertexArray>();
      primitive_helper::BuildRectXY(2.f, *vertex_array);
      renderable->AddVertexArray(vertex_array,
        fmt::format("mat:msi_{:03d}", i).c_str());
      renderable->SetQueueId(Renderable::kQueueTransparent + total_layers - i);
      renderable->SetEffect("effect:msi");
      scene->Renderables().push_back(renderable);
    }

    // shader material, effect and render pass
    std::vector<std::string> basis_k_names;
    std::vector<std::string> basis_cubemap_names;
    for (int i = 0; i < basis_out; ++i) {
      basis_k_names.emplace_back(fmt::format("texture:basis_k_{:02d}", i));
      basis_cubemap_names.emplace_back(fmt::format("texture:basis_cubemap_{:02d}", i));
    }
    for (int i = 0; i < total_layers; ++i) {
      auto mat = std::make_shared<Material>();
      mat->SetUniformValue("theta_phi_range",
        glm::vec2(theta_half_range * 2.f, phi_half_range * 2.f));
      mat->SetUniformValue("radius", radii[i]);
      mat->SetUniformValue("layer_id", i);
      mat->SetUniformValue("msi_w", msi_w);
      mat->SetUniformValue("msi_h", msi_h);
      mat->SetUniformValue("msi_layers", msi_layers);
      mat->SetUniformValue("msi_sublayers", msi_sublayers);
      mat->SetTexture("msi_rgb", {"texture:msi_rgb"});
      mat->SetTexture("msi_alpha", {"texture:msi_alpha"});
      mat->SetTexture("basis_k[0]", basis_k_names);
      mat->SetTexture("basis_cubemap[0]", basis_cubemap_names);
      en.ResrcMgr().Add(fmt::format("mat:msi_{:03d}", i), mat);
    }

    std::vector<std::unique_ptr<RenderState>> render_states;
    render_states.push_back(std::make_unique<BlendEnableState>(true));
    render_states.push_back(std::make_unique<BlendFuncState>(
      render_state::kBlendSrcAlpha, render_state::kBlendOneMinusSrcAlpha));
    CreateEffectFromFileHelper("effect:msi", "msi.vs", "msi.fs", nullptr,
      std::move(render_states));

    RenderPass render_pass;
    en.RenderPasses().push_back(render_pass);

    // load data
    cnpy::NpyArray msi_rgb = cnpy::npy_load(folder_ / "msi_rgb.npy");
    cnpy::NpyArray msi_alpha = cnpy::npy_load(folder_ / "msi_alpha.npy");
    std::vector<cnpy::NpyArray> basis_k;
    std::vector<cnpy::NpyArray> basis_cubemap;
    for (int i = 0; i < basis_out; ++i) {
      basis_k.emplace_back(std::move(cnpy::npy_load(
        folder_ / fmt::format("basis_k_{:02d}.npy", i))));
      basis_cubemap.emplace_back(std::move(cnpy::npy_load(
        folder_ / fmt::format("basis_{:02d}.npy", i))));
    }
    int cubemap_size = basis_cubemap.back().shape.back();
    print_npy_array_info(msi_rgb);
    print_npy_array_info(msi_alpha);
    print_npy_array_info(basis_k.back());
    print_npy_array_info(basis_cubemap.back());

    // msi rgb texture data
    int bpc = 32;
    int pitch = msi_w * 3 * bpc / 8;
    auto msi_rgb_roi = imgpp::ImgROI((uint8_t*)msi_rgb.data_holder->data(),
      msi_w, msi_h, msi_layers, 3, bpc, pitch, pitch * msi_h, true, true);
    imgpp::Img msi_rgb_img;
    msi_rgb_img.CopyFrom(msi_rgb_roi);
    auto msi_rgb_src = texture_helper::CreateTextureSrc(msi_rgb_img);
    TextureDesc msi_rgb_desc;
    texture_helper::CreateTextureDesc(msi_rgb_src, false, false,
      TextureDesc::kLinearMipmapLinear, TextureDesc::kLinear,
      TextureDesc::kClampToBorder, TextureDesc::kClampToBorder, 3, msi_rgb_desc);
    texture_helper::CreateTextureFromDesc("texture:msi_rgb", msi_rgb_desc);

    // msi alpha texture data
    bpc = 32;
    pitch = msi_w * bpc / 8;
    auto msi_alpha_roi = imgpp::ImgROI((uint8_t*)msi_alpha.data_holder->data(),
      msi_w, msi_h, total_layers, 1, bpc, pitch, pitch * msi_h, true, true);
    imgpp::Img msi_alpha_img;
    msi_alpha_img.CopyFrom(msi_alpha_roi);
    auto msi_alpha_src = texture_helper::CreateTextureSrc(msi_alpha_img);
    TextureDesc msi_alpha_desc;
    texture_helper::CreateTextureDesc(msi_alpha_src, false, false,
      TextureDesc::kLinearMipmapLinear, TextureDesc::kLinear,
      TextureDesc::kClampToBorder, TextureDesc::kClampToBorder, 3, msi_alpha_desc);
    texture_helper::CreateTextureFromDesc("texture:msi_alpha", msi_alpha_desc);

    // basis k texture data
    bpc = 16;
    pitch = msi_w * 3 * bpc / 8;
    for (int i = 0; i < basis_out; ++i) {
      auto basis_k_roi = imgpp::ImgROI((uint8_t*)basis_k[i].data_holder->data(),
        msi_w, msi_h, msi_layers, 3, bpc, pitch, pitch * msi_h, true, true);
      imgpp::Img basis_k_img;
      basis_k_img.CopyFrom(basis_k_roi);
      auto basis_k_src = texture_helper::CreateTextureSrc(basis_k_img);
      TextureDesc basis_k_desc;
      texture_helper::CreateTextureDesc(basis_k_src, false, false,
        TextureDesc::kLinearMipmapLinear, TextureDesc::kLinear,
        TextureDesc::kClampToBorder, TextureDesc::kClampToBorder, 3, basis_k_desc);
      texture_helper::CreateTextureFromDesc(
        fmt::format("texture:basis_k_{:02d}", i).c_str(), basis_k_desc);
    }

    // basis cubemap texture data
    bpc = 16;
    pitch = cubemap_size * bpc / 8;
    for (int i = 0; i < basis_out; ++i) {
      auto basis_cubemap_src = std::make_shared<ImgppTextureSrc>(
        6, 1, 1, imgpp::FORMAT_UNDEFINED);
      auto cubemap_roi = imgpp::ImgROI((uint8_t*)basis_cubemap[i].data_holder->data(),
        cubemap_size, cubemap_size, 6, 1, bpc, pitch, pitch * cubemap_size, true, true);
      imgpp::Img cubemap_img;
      cubemap_img.CopyFrom(cubemap_roi);
      basis_cubemap_src->AddBuffer(cubemap_img.Data());
      basis_cubemap_src->AddROI(cubemap_img.ROI());
      TextureDesc basis_cubemap_desc;
      texture_helper::CreateTextureDesc(basis_cubemap_src, false, false,
        TextureDesc::kLinearMipmapLinear, TextureDesc::kLinear,
        TextureDesc::kClampToEdge, TextureDesc::kClampToEdge,
        2, basis_cubemap_desc);
      std::shared_ptr<TextureCubeMap> texture(new TextureCubeMap);
      if (texture->Create(basis_cubemap_desc)) {
        eustis::Engine::Instance().ResrcMgr().Add(
          fmt::format("texture:basis_cubemap_{:02d}", i), texture);
      }
    }
  }

private:
  bool flip_y_{false};
  std::filesystem::path folder_;
  nlohmann::json params_;
  std::shared_ptr<FPSController> cam_ctrl_;
};

}

int main(int argc, char **argv) {
  if (argc < 2) {
    std::cout << "Usage: " << argv[0] << " data_folder [flip_y]" << std::endl;
    return 0;
  }

  bool flip_y = false;
  if (argc >= 3 && atoi(argv[2]) != 0) {
    flip_y = true;
  }

  AppFrame app;
  std::shared_ptr<MsiEntity> entity(new MsiEntity(argv[1], flip_y));
  Engine::Instance().EntityMgr().Add("msi_entity", entity);
  app.Run(800, 600, std::nullopt);

  return 0;
}
